from cryptography.fernet import Fernet
import os
import argparse

textchars = bytearray({7,8,9,10,12,13,27} | set(range(0x20, 0x100)) - {0x7f})
is_binary_string = lambda bytes: bool(bytes.translate(None, textchars))

def is_binary(filename):
    with open(filename, 'rb') as f:
        return is_binary_string(f.read(1024))

def find_all_files(base):
    for root, dirs, fs in os.walk(base):
        for f in fs:
            fullname = os.path.join(root, f)
            relpath = os.path.relpath(fullname, base)
            # print(f'root: {root}, dirs: {dirs}, f: {f}, relpath: {relpath}')
            yield relpath

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--input_dir',
        default="",
        required=True,
        help='dir to be encrypted'
    )
    parser.add_argument(
        '--output_dir',
        default='',
        required=True,
        help='output dir'
    )
    parser.add_argument(
        '--key',
        default='',
        required=True,
        help='encrypt key'
    )

    args, rest_args = parser.parse_known_args()
    print(args)
    input_dir = args.input_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    print(f'input_dir: {input_dir}, output_dir: {output_dir}')

    key = args.key
    if len(key) == 0:
        print('encrypt key must not be empty')
        exit(1)

    fernet = Fernet(key)

    with open('file_map.csv', 'w') as conf:
        for f in find_all_files(input_dir):
            print(f'output_dir: {output_dir}')
            outfile = os.path.join(output_dir, f)
            infile = os.path.join(input_dir, f)
            print(f'input: {infile}, output: {outfile}')
            cur_dir = os.path.dirname(outfile)
            if not os.path.exists(cur_dir):
                os.makedirs(cur_dir)
            is_bin = is_binary(infile)
            bin_mode = 'b' if is_bin else ''
            #with open(outfile, f'w{bin_mode}') as ew:
            with open(outfile, f'wb') as ew:
                with open(infile, f'rb') as src:
                    content = src.read()
                    if not isinstance(content, bytes):
                        content = content.encode('utf-8')
                    encrypted_content = fernet.encrypt(content)
                    #if not is_bin:
                    #    encrypted_content = encrypted_content.decode('utf-8')
                    print(f'raw len: {len(content)}, encrypted len: {len(encrypted_content)}')
                    # decrypt_content = fernet.decrypt(encrypted_content)
                    # assert content == decrypt_content
                    ew.write(encrypted_content)
            conf.write(f'{f},{outfile}\n')



